from os.path import exists,join,basename
from os import mkdir,remove
from six.moves import urllib
import tarfile
from torchvision.transforms import Compose,CenterCrop,ToTensor,Scale,Resize
from my_data_utils import My_datasets


def input_transform(crop_size,upscale_factor):
    return Compose([CenterCrop(crop_size),
                    Resize((crop_size//upscale_factor, crop_size//upscale_factor)),
                    Resize((crop_size,crop_size)),
                    ToTensor(),])

def target_tranform(crop_size):
    return Compose([CenterCrop(crop_size), ToTensor(),])

def calculate_valid_crop_size(crop_size,upscale_factor):
    return crop_size - (crop_size % upscale_factor)

def get_trainning_set(upscale_factor):
    root_dir = './BSDS300/images'
    train_dir = join(root_dir,'train')
    crop_size = calculate_valid_crop_size(256,upscale_factor)
    return My_datasets(train_dir,input_transform=input_transform(crop_size,upscale_factor),
                       target_transform = target_tranform(crop_size))

def get_test_set(upscale_factor):
    root_dir = './BSDS300/images'
    test_dir = join(root_dir,'test')
    crop_size = calculate_valid_crop_size(256,upscale_factor)
    return My_datasets(test_dir,input_transform=input_transform(crop_size,upscale_factor),
                        target_transform = target_tranform(crop_size))






